import os
import re
import sys
import glob
import shutil

root = os.path.dirname(os.path.dirname(__file__))
sys.path.append(root)
from util_my import add_arg, replace_content
from project.yolov5.data import refactor

def export_formats():
    import pandas as pd
    x = [
        ['PyTorch', '-', '.pt', True, True],
        ['TorchScript', 'torchscript', '.torchscript', True, True],
        ['ONNX', 'onnx', '.onnx', True, True],
        ['OpenVINO', 'openvino', '_openvino_model', True, False],
        ['TensorRT', 'engine', '.engine', False, True],
        ['CoreML', 'coreml', '.mlmodel', True, False],
        ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],
        ['TensorFlow GraphDef', 'pb', '.pb', True, True],
        ['TensorFlow Lite', 'tflite', '.tflite', True, False],
        ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
        ['TensorFlow.js', 'tfjs', '_web_model', False, False],
        ['PaddlePaddle', 'paddle', '_paddle_model', True, True],]
    return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])

def run(task='', data_dir='', cfg_id='', github_dir='', device='', weights='', out_dir='', 
    frameworks='', logger=None):

    from project.yolov5.infer import get_weight
    git_dir, code_dir = refactor(task=task, data_dir=data_dir, github_dir=github_dir)
    weights = get_weight(git_dir=git_dir, cfg_id=cfg_id, weights=weights, out_dir=out_dir)

    frameworks = frameworks.replace('tensorrt', 'engine').split(',') if type(frameworks) is str else frameworks
    add_arg('--device', device)
    add_arg('--weights', weights)

    torchscript = '''
    input_shapes = [[1, 3, 640, 640]]
    output_shapes = []
    out_ = model(im)

    for out in out_:
        if type(out) == list:
            for o in out:
                output_shapes.append(list(o.shape))
        else:
            output_shapes.append(list(out.shape))

    shapes = {
        'input_shapes': input_shapes,
        'output_shapes': output_shapes,
    }
    extra_files = {
        'shapes': json.dumps(shapes)
    }
    '''

    def replace_export(content):
        if torchscript not in content:
            content = re.sub('(\n' + ' ' * 4 + 'if optimize:  # )', torchscript + r'\1', content)
        return content

    replace_content(os.path.join(code_dir, 'export.py'), replace_export)

    root = os.path.dirname(os.path.dirname(__file__))
    while root in sys.path:
        sys.path.remove(root)

    try:
        import downloads.yolov5.export as export
    except:
        sys.path.append(os.path.dirname(git_dir))
        import yolov5.export as export

    opt = export.parse_opt()
    opt.weights = weights
    opt.include = frameworks
    try:
        export.run(**vars(opt))
    except Exception as e:
        if logger:
            logger.log(logger.ERROR, e, exception=True)

    for weights in glob.glob(re.sub('.\w+$', '**', weights)):
        if '.pt' in weights: continue
        os.makedirs(out_dir, exist_ok=True)
        shutil.move(weights, out_dir)
